#include <algorithm>
#include <cassert>
#include <chrono>
#include "BVH.hpp"

BVHAccel::BVHAccel(std::vector<Object*> p, int maxPrimsInNode,
                   SplitMethod splitMethod)
    : maxPrimsInNode(std::min(255, maxPrimsInNode)), splitMethod(splitMethod),
      primitives(std::move(p))
{
    std::chrono::system_clock::time_point start = std::chrono::system_clock::now();
    if (primitives.empty())
        return;
    if(splitMethod == SplitMethod::NAIVE){
        root = recursiveBuild(primitives);
    }
    else{
        root = recursiveBuild2(primitives);
    }

    std::chrono::system_clock::time_point stop = std::chrono::system_clock::now();
    std::chrono::seconds s = std::chrono::duration_cast<std::chrono::seconds>(stop - start);
    std::chrono::microseconds m(std::chrono::duration_cast<std::chrono::microseconds>(stop - start).count() % 1000000);

    printf(
        "\rBVH Generation complete: \nTime Taken: %ld seconds, %ld micro seconds\n\n", s.count(), m.count());
}

BVHBuildNode* BVHAccel::recursiveBuild(std::vector<Object*>& objects)
{
    BVHBuildNode* node = new BVHBuildNode();

    // Compute bounds of all primitives in BVH node
    Bounds3 bounds;
    for (int i = 0; i < objects.size(); ++i)
        bounds = Union(bounds, objects[i]->getBounds());
    if (objects.size() == 1) {
        // Create leaf _BVHBuildNode_
        node->bounds = objects[0]->getBounds();
        node->object = objects[0];
        node->left = nullptr;
        node->right = nullptr;
        return node;
    }
    else if (objects.size() == 2) {
        auto tmp1 = std::vector<Object*>{objects[0]}, tmp2 = std::vector<Object*>{objects[1]};
        node->left = recursiveBuild(tmp1);
        node->right = recursiveBuild(tmp2);

        node->bounds = Union(node->left->bounds, node->right->bounds);
        return node;
    }
    else {
        Bounds3 centroidBounds;
        for (int i = 0; i < objects.size(); ++i)
            centroidBounds =
                Union(centroidBounds, objects[i]->getBounds().Centroid());
        int dim = centroidBounds.maxExtent();
        switch (dim) {
        case 0:
            std::sort(objects.begin(), objects.end(), [](auto f1, auto f2) {
                return f1->getBounds().Centroid().x <
                       f2->getBounds().Centroid().x;
            });
            break;
        case 1:
            std::sort(objects.begin(), objects.end(), [](auto f1, auto f2) {
                return f1->getBounds().Centroid().y <
                       f2->getBounds().Centroid().y;
            });
            break;
        case 2:
            std::sort(objects.begin(), objects.end(), [](auto f1, auto f2) {
                return f1->getBounds().Centroid().z <
                       f2->getBounds().Centroid().z;
            });
            break;
        }

        auto beginning = objects.begin();
        auto middling = objects.begin() + (objects.size() / 2);
        auto ending = objects.end();

        auto leftshapes = std::vector<Object*>(beginning, middling);
        auto rightshapes = std::vector<Object*>(middling, ending);

        assert(objects.size() == (leftshapes.size() + rightshapes.size()));

        node->left = recursiveBuild(leftshapes);
        node->right = recursiveBuild(rightshapes);

        node->bounds = Union(node->left->bounds, node->right->bounds);
    }

    return node;
}

BVHBuildNode* BVHAccel::recursiveBuild2(std::vector<Object*>& objects){
    int bucket = 24, obj_num = objects.size() / bucket;
    obj_num = std::max(1, obj_num);

    BVHBuildNode* node = new BVHBuildNode();
    if(objects.size() == 1){
        node->bounds = objects[0]->getBounds();
        node->object = objects[0];
        node->left = nullptr, node->right = nullptr;
        return node;
    }
    else if(objects.size() == 2){
        auto tmp1 = std::vector<Object*>{objects[0]}, tmp2 = std::vector<Object*>{objects[1]};
        node->left = recursiveBuild2(tmp1);
        node->right = recursiveBuild2(tmp2);
        node->bounds = Union(node->left->bounds, node->right->bounds);
        return node; 
    }

    auto cal_func = [](const std::vector<Object*>& obj1, const std::vector<Object*>& obj2){
        Bounds3 b1, b2;
        int s1 = obj1.size(), s2 = obj2.size();
        for(auto j: obj1){
            b1 = Union(b1, j->getBounds());
        }
        for(auto j: obj2){
            b2 = Union(b2, j->getBounds());
        }  
        return s1 * std::abs(b1.Diagonal().x * b1.Diagonal().y * b1.Diagonal().z) + \
        s2 * std::abs(b2.Diagonal().x * b2.Diagonal().y * b2.Diagonal().z);       
    };

    auto sort_func = [obj_num, &objects, &cal_func](int dim, int sig) -> std::tuple<float, int>{
        switch (dim)
        {
        case 0:
            std::sort(objects.begin(), objects.end(), [](auto f1, auto f2) {
                return f1->getBounds().Centroid().x <
                       f2->getBounds().Centroid().x;
            });
            break;
        case 1:
            std::sort(objects.begin(), objects.end(), [](auto f1, auto f2) {
                return f1->getBounds().Centroid().y <
                       f2->getBounds().Centroid().y;
            });
            break;
        case 2:
            std::sort(objects.begin(), objects.end(), [](auto f1, auto f2) {
                return f1->getBounds().Centroid().z <
                       f2->getBounds().Centroid().z;
            });
            break;
        }
        if(sig == 1){
            return std::tuple<float, int>(0, 0);
        }     
        auto mid = objects.begin() + obj_num;
        float rel = std::numeric_limits<float>::max();
        int parse_num = -1;
        for(;mid < objects.end(); mid += obj_num){
            float tmp = cal_func(std::vector<Object*>{objects.begin(), mid}, std::vector<Object*>{mid, objects.end()}); 
            if(tmp < rel){
                parse_num = mid - objects.begin();
                rel = tmp;
            }        
        }
        return std::tuple<float, int>(rel, parse_num);                                
    };
    
    std::array<std::tuple<float, int>, 3> tmp;
    for(int i=0; i < 3; ++i){
        tmp[i] = sort_func(i, 0);
    }
    auto ptr = std::min_element(tmp.begin(), tmp.end());
    sort_func(ptr - tmp.begin(), 1);

    std::vector<Object*> vec1{objects.begin(), objects.begin() + std::get<1>(*ptr)};
    std::vector<Object*> vec2{objects.begin() + std::get<1>(*ptr), objects.end()};
    node->left = recursiveBuild2(vec1);
    node->right = recursiveBuild2(vec2);
    node->bounds = Union(node->left->bounds, node->right->bounds);       
    return node;   
}

Intersection BVHAccel::Intersect(const Ray& ray) const
{
    Intersection isect;
    if (!root)
        return isect;
    isect = BVHAccel::getIntersection(root, ray);
    return isect;
}

Intersection BVHAccel::getIntersection(BVHBuildNode* node, const Ray& ray) const
{
    // TODO Traverse the BVH to find intersection
    if(node->bounds.IntersectP(ray)){
        Intersection left, right;
        if(!(node->left || node->right)){
            return node->object->getIntersection(ray);
        }
        if(node->left){
            left = getIntersection(node->left,ray);            
        }
        if(node->right){
            right = getIntersection(node->right, ray);   
        }
        return std::min(left, right, [](const Intersection& a, const Intersection& b){
            return a.distance < b.distance;
        });             
    }
    return Intersection(); 
}